from itertools import izip

from opt.solver import Solver
from opt.treesearch.node import TreeNode

from .metanode import MCTS_MetaNode
from .params import MCTS_ParamSet
from .treeexhausted import TreeExhausted


class MCTS(Solver):
    """
    A generic Monte-Carlo Tree Search solver. In most cases, this should work seamlessly with
    existing TreeNode subclasses. For more advanced usage, MCTS_MetaNode can be subclassed and
    used (through the 'metanode_factory' parameter) instead of the default meta-node class.
    """
    ParamSet = MCTS_ParamSet          # paramset class used by this solver
    node_factory = None               # function used to create the root node
    metanode_factory = MCTS_MetaNode  # function used to create metanodes
    default_plugins = list(Solver.default_plugins) + [TreeExhausted]

    def _extract_solution_and_meta(self, node):
        return node.solution_and_meta()

    def _bootstrap(self):
        """Bootstrap the search by creating the root node and running a simulation from it."""
        root_node = self.params.node_factory(self)
        root_meta = self.params.metanode_factory(self, root_node)
        self.root = root_meta
        sim_results = self._simulation_step([root_meta])
        self._backpropagation_step([root_meta], sim_results)

    def _iterate(self):
        """An iteration of MCTS consists of four steps:
            1) selection of a non-fully expanded node in the tree
            2) expansion of one or more children of the selected node
            3) for each new node, run a random simulation until a terminal state is reached
            4) use the results of the simulations to update the evaluations of the selected node
               and its ancestors
        """
        selected = self._selection_step()
        expanded = self._expansion_step(selected)
        sim_results = self._simulation_step(expanded)
        self._backpropagation_step(expanded, sim_results)

    def _selection_step(self):
        """Descend through the tree until we find an expandable node (i.e. a node which hasn't
        been fully expanded yet).  This step uses a selection policy that, given a list of
        metanodes, should return a list of candidates that are considered best according to some
        criterion.  The default criterion is to maximize the score in a UCT-like formula."""
        selection_policy = self.params.selection_policy
        random_choice = self.rng.choice
        metanode = self.root
        while not metanode.is_expandable:
            if len(metanode.children) == 1:
                metanode = metanode.children[0]
            else:
                candidates = selection_policy(metanode.children)
                if len(candidates) == 1:
                    metanode = candidates[0]
                else:
                    metanode = random_choice(candidates)
        return metanode

    def _expansion_step(self, selected):
        expansion_policy = self.params.expansion_policy
        expanded = []
        while not selected.is_fully_expanded and len(expanded) < self.params.expansion_limit:
            child = expansion_policy(selected)
            if child.is_fully_expanded:
                self.solutions.check(child.node)
            else:
                expanded.append(child)
        return expanded

    def _simulation_step(self, expanded):
        sim_policy = self.param.simulation_policy
        sim_results = []
        for metanode in expanded:
            sim_result = sim_policy(metanode)
            sim_results.append(sim_result)
        return sim_results

    def _backpropagation_step(self, expanded, sim_results):
        backpropagation_policy = self.params.backpropagation_policy
        for metanode, sim_result in izip(expanded, sim_results):
            if isinstance(sim_result, TreeNode):
                self.solutions.check(sim_result)
                sim_result = sim_result.objective()
            else:
                backpropagation_policy(metanode, sim_result)

    def _prune_tree(self):
        pruned_nodes = 0
        UB = self.upper_bound
        stack = [self.root]
        while len(stack) > 0:
            meta = stack.pop()
            if meta.node.lower_bound() >= UB:
                pruned_nodes += meta.dimensions[0]
                meta.chop()
            elif len(meta.children) > 0:
                stack.extend(meta.children)
        self.log.debug("PRUNE TREE: removed {} nodes".format(pruned_nodes))
